package com.thoughtful.registry.zk.util;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import com.thougtful.enums.RpcConfigEnum;
import com.thougtful.utils.PropertiesFileUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;

/**
 * @author yushun
 * @createTime 2022/10/19 - 14:40
 **/
@Slf4j
public final class CuratorUtils {

    private static final int BASE_SLEEP_TIME = 1000;

    private static final int MAX_RETRIES = 3;

    public static final String ZK_REGISTER_ROOT_PATH = "/rpc";

    /**
     * 存储服务和服务器地址列表的映射关系缓存
     * <p>
     * k: 服务地址
     * v: 服务对应服务器地址
     */
    private static final Map<String, List<String>> SERVICE_ADDRESS_MAP = new ConcurrentHashMap<>();

    /**
     * 存储已经注册过的服务列表
     */
    private static final Set<String> REGISTERED_PATH_SET = ConcurrentHashMap.newKeySet();

    private static CuratorFramework zkClient;

    public static final String DEFAULT_ZOOKEEPER_ADDRESS = "127.0.0.1:2181";

    private CuratorUtils() {}

    public static CuratorFramework getZkClient() {
        // 检测用户是否自定义了 ZK 的地址
        Properties properties = PropertiesFileUtil.readPropertiesFile(RpcConfigEnum.RPC_CONFIG_PATH.getPropertyValue());
        String zookeeperAddress = Optional.ofNullable(properties)
            .map(t -> t.getProperty(RpcConfigEnum.ZK_ADDRESS.getPropertyValue(), DEFAULT_ZOOKEEPER_ADDRESS))
            .orElse(DEFAULT_ZOOKEEPER_ADDRESS);

        // 如果已有 ZK 客户端, 并且已启动, 则直接返回
        if (Objects.nonNull(zkClient) && zkClient.getState() == CuratorFrameworkState.STARTED) {
            return zkClient;
        }
        // 定义重试策略, 重试三次, 并且每次重试间隔 1000ms
        ExponentialBackoffRetry retryPolicy = new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES);
        zkClient = CuratorFrameworkFactory.builder()
            .connectString(zookeeperAddress)
            .retryPolicy(retryPolicy)
            .build();
        zkClient.start();
        try {
            // 等待 30s 连接 Zookeeper
            if (!zkClient.blockUntilConnected(30, TimeUnit.SECONDS)) {
                throw new RuntimeException("Time out waiting to connect to ZK!");
            }
        } catch (InterruptedException e) {
            log.error("zkClient start error!", e);
        }
        return zkClient;
    }

    /**
     * 创建一个持久的 ZK 节点, 客户端断开连接时, 持久节点也不会被删除
     */
    public static void createPersistentNode(CuratorFramework zkClient, String servicePath) {
        try {
            if (REGISTERED_PATH_SET.contains(servicePath) || zkClient.checkExists().forPath(servicePath) != null) {
                log.info("The node already exists, The node is: [{}]", servicePath);
            } else {
                zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(servicePath);
                log.info("The node was created successfully, The node is: [{}]", servicePath);
            }
        } catch (Exception e) {
            log.error("create persistent node for path [{}] fail", servicePath);
        }
    }

    public static List<String> getChildrenNodes(CuratorFramework zkClient, String rpcServiceName) {
        if (SERVICE_ADDRESS_MAP.containsKey(rpcServiceName)) {
            return SERVICE_ADDRESS_MAP.get(rpcServiceName);
        }
        List<String> result = null;
        String servicePath = ZK_REGISTER_ROOT_PATH + "/" +rpcServiceName;
        try {
            result = zkClient.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, result);
            registerWatcher(rpcServiceName, zkClient);
        } catch (Exception e) {
            log.error("get children nodes for path [{}] fail", servicePath);
        }
        return result;
    }

    private static void registerWatcher(String rpcServiceName, CuratorFramework zkClient) throws Exception {
        String servicePath = ZK_REGISTER_ROOT_PATH + "/" + rpcServiceName;
        PathChildrenCache pathChildrenCache = new PathChildrenCache(zkClient, servicePath, true);
        PathChildrenCacheListener pathChildrenCacheListener = (curatorFramework, pathChildrenCacheEvent) -> {
            List<String> serviceAddress = curatorFramework.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, serviceAddress);
        };
        pathChildrenCache.getListenable().addListener(pathChildrenCacheListener);
        pathChildrenCache.start();
    }
}
